package cn.liu.knowledge.advisor;

import com.alibaba.fastjson.JSONObject;
import lombok.extern.slf4j.Slf4j;
import netscape.javascript.JSObject;
import org.reactivestreams.Publisher;
import org.springframework.ai.chat.client.advisor.api.AdvisedRequest;
import org.springframework.ai.chat.client.advisor.api.AdvisedResponse;
import org.springframework.ai.chat.client.advisor.api.StreamAroundAdvisor;
import org.springframework.ai.chat.client.advisor.api.StreamAroundAdvisorChain;
import org.springframework.ai.chat.model.MessageAggregator;
import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.core.io.buffer.DataBufferFactory;
import org.springframework.core.io.buffer.DefaultDataBufferFactory;
import org.springframework.http.client.reactive.ClientHttpRequest;
import org.springframework.http.codec.HttpMessageWriter;
import org.springframework.http.server.reactive.ServerHttpRequest;
import org.springframework.web.reactive.function.BodyInserter;
import org.springframework.web.reactive.function.BodyInserters;
import org.springframework.web.reactive.function.client.*;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.util.List;
import java.util.Map;
import java.util.Optional;

@Slf4j
public class RagAdvisor implements StreamAroundAdvisor {

    private String prompt;

    public RagAdvisor(String prompt) {
        this.prompt = prompt;
    }

    @Override
    public Flux<AdvisedResponse> aroundStream(AdvisedRequest advisedRequest, StreamAroundAdvisorChain chain) {

        log.debug("BEFORE: {}", advisedRequest);
//        OllamaOptions.builder().model("model").build();
        Flux<AdvisedResponse> advisedResponses = chain.nextAroundStream(advisedRequest);

        Flux<AdvisedResponse> advisedResponseFlux = new MessageAggregator().aggregateAdvisedResponse(advisedResponses,
                advisedResponse -> log.debug("AFTER: {}", advisedResponse));

        return advisedResponseFlux;
    }

    @Override
    public String getName() {
        return this.getClass().getName();
    }

    @Override
    public int getOrder() {
        return Integer.MAX_VALUE;
    }


    public static void main(String[] args) {
        WebClient client = WebClient.builder()
//                .exchangeFunction(new ExchangeFunction() {
//                    @Override
//                    public Mono<ClientResponse> exchange(ClientRequest request) {
//                        System.out.println(request.body());
////                        return null;
//                        return Mono.just(request);
//                    }
//                })
                .filter(logRequest())
                .filter(logResponse())
                .build();

        JSONObject param = new JSONObject();
        param.put("model","qwen2.5:3b");
        param.put("prompt","你好");
        param.put("stream",false);
        Mono<JSONObject> just = Mono.just(param);
        client.post()
                .uri("http://localhost:11434/api/generate")
                .bodyValue(param)
                .retrieve()
                .bodyToMono(String.class)
                .subscribe(System.out::println);
    }

    public static ExchangeFilterFunction logRequest() {
        return ExchangeFilterFunction.ofRequestProcessor(clientRequest -> {
            log.debug("Request: " + clientRequest.method() + " " + clientRequest.url());



            clientRequest.headers().forEach((name, values) ->
                    values.forEach(value -> System.out.println(name + ": " + value)));

            BodyInserter<?, ? super ClientHttpRequest> body = clientRequest.body();
            if (clientRequest.body() instanceof BodyInserter<?, ?>) {
                // 根据不同类型处理
                if (body instanceof BodyInserters.FormInserter) {
                    BodyInserters.FormInserter<?> formInserter = (BodyInserters.FormInserter<?>) body;
                    System.out.println("Form data: " + formInserter);
                } else if (body instanceof BodyInserters.MultipartInserter) {
                    BodyInserters.MultipartInserter multipartInserter = (BodyInserters.MultipartInserter) body;
                    System.out.println("Multipart data: " + multipartInserter);
                }
            }

            log.debug(clientRequest.body().toString());
            return Mono.just(clientRequest);
        });
    }

    public static ExchangeFilterFunction logResponse() {
        return ExchangeFilterFunction.ofResponseProcessor(clientResponse -> {
            System.out.println("Response status: " + clientResponse.statusCode());
            clientResponse.headers().asHttpHeaders().forEach((name, values) ->
                    values.forEach(value -> System.out.println(name + ": " + value)));
            return Mono.just(clientResponse);
        });
    }
}

